import random
from typing import Dict, List, Tuple
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Normal
import socket
import json
import os
import pickle

def init_layer_uniform(layer: nn.Linear, init_w: float = 3e-3) -> nn.Linear:
    """Init uniform parameters on the single layer."""
    layer.weight.data.uniform_(-init_w, init_w)
    layer.bias.data.uniform_(-init_w, init_w)

    return layer

class Actor(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, hidden_num: int =64, init_w: float = 3e-3):
        """Initialize."""
        super(Actor, self).__init__()

        self.hidden1 = nn.Linear(in_dim, hidden_num)
        self.out = nn.Linear(hidden_num, out_dim)

        self.out.weight.data.uniform_(-init_w, init_w)
        self.out.bias.data.uniform_(-init_w, init_w)

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        """Forward method implementation."""
        x = F.relu(self.hidden1(state))
        action = self.out(x).tanh()

        return action


def decode_data(data):
    # load json
    recv_obser = json.loads(data)
    # print(recv_obser)
    observation = []
    observation_dict = recv_obser['observation']
    for key in observation_dict:
        observation.append(observation_dict[key])
    reward = recv_obser['reward']
    terminated = recv_obser['terminated']
    truncated = recv_obser['truncated']
    return observation, reward, terminated, truncated

def encode_data(action, reset_flag=1):
    rudder = float(action) * 5
    print(f'rudder is {rudder}')
    if rudder >= 30.0:
        rudder = 30.0
    elif rudder <= -30.0:
        rudder = -30.0
    
    origin_data = {'boatname':'SLM7001',
                   'restart': reset_flag,
                   'rudl': rudder,
                   'rudr': rudder,
                   'rspl': 800,
                   'rspr': 800,
                   'subSystem': "control"
                   }

    data = json.dumps(origin_data, sort_keys=True, indent=4, separators=(',', ':'))
    return data.encode('utf-8')

hidden_num = 64
if __name__ == '__main__':
    # local socket
    # LOCAL_HOST = '192.168.0.137'
    LOCAL_HOST = '127.0.0.1'
    LOCAL_PORT = 10901
    # network socket
    REMOTE_HOST = '127.0.0.1'
    # REMOTE_HOST = '192.168.0.191'
    REMOTE_PORT = 10905
    udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_socket.bind((LOCAL_HOST, LOCAL_PORT))
    
    obs_dim = 3
    action_dim =1
    
    HOME = os.getcwd()
    print(HOME)
    model_path = os.path.join(HOME, 'TD3/TD3_model.pt')
    model = torch.load(model_path)['policy_net_state_dict']
    
    # device: cpu / gpu
    device = torch.device(
        "cuda" if torch.cuda.is_available() else "cpu"
    )
    print(device)
    actor = Actor(obs_dim, action_dim,hidden_num=hidden_num).to(device=device)
    
    actor.load_state_dict(model)
    
    # reset
    action_data = encode_data(0, reset_flag=0)
    udp_socket.sendto(action_data, (REMOTE_HOST, REMOTE_PORT))# action
    
    mean_time = 0
    count = 0
    while(1):
        count += 1
        time_start = time.time()
        # Initialize the environment and get its state
        info, addr = udp_socket.recvfrom(1024)
        observation, _, _, _= decode_data(info)
        state = torch.tensor(observation, dtype=torch.float32, device=device).unsqueeze(0)
        
        selected_action = actor(
            torch.FloatTensor(state).to(device)
        )[0].detach().cpu().numpy()
        
        print(f'select action is: {selected_action}')
        
        action_data = encode_data(selected_action, reset_flag=0)
        udp_socket.sendto(action_data, (REMOTE_HOST, REMOTE_PORT))# action
        
        time_end = time.time()
        consume_time = (time_end - time_start) * 1000.0
        
        mean_time = mean_time + consume_time
        
        print(f"consume time is {mean_time/count} ms")
    udp_socket.close()